//
// Created by jizhe hu on 2024/7/10.
//

#ifndef SYNC_SET_HPP
#define SYNC_SET_HPP

#include <vector>
#include <random>
#include <iostream>

namespace sync {

template<class ElemType>
class BaseSet {
public:
    virtual void add(ElemType elem) = 0;
    virtual bool find(ElemType elem) = 0;
    virtual bool erase(ElemType elem) = 0;
};

template<class ElemType, class RWMutexType, class DeleteMutexType>
class Set : public BaseSet<ElemType> {
private:
    struct Node {
        ElemType key;
        std::vector<Node*> forwards;

        Node(ElemType key, int level) : key(key) {
            forwards.resize(level, nullptr);
        }
    };

private:
    // 最大层数
    int m_max_level;
    // 元素个数
    size_t m_sz;
    // 概率值
    double m_prob;
    // 头节点
    Node* m_head;

    std::random_device m_rd;
    std::mt19937 m_gen;
    std::uniform_real_distribution<double> m_dis;

    // 生成随机层数
    int randomLevel() {
        int level = 1;
        while (m_dis(m_gen) < m_prob && level < m_max_level) {
            level++;
        }
        return level;
    }

public:
    Set(double p = 0.5, int max_level = 32) : m_gen(m_rd()), m_dis(0.0, 1.0), m_prob(p), m_max_level(max_level), m_sz(0) {
        m_head = new Node(INT_MIN, m_max_level);
    }

    ~Set() {
        Node* curr = m_head;
        while (curr) {
            Node* tmp = curr;
            curr = curr->forwards[0];
            delete tmp;
        }
    }

    bool find(ElemType target) {
        Node* curr = m_head;
        for (int i = m_max_level - 1; i >= 0; i--) {
            while (curr->forwards[i] && curr->forwards[i]->key < target) {
                curr = curr->forwards[i];
            }
        }
        curr = curr->forwards[0];
        return curr && curr->key == target;
    }

    void add(ElemType elem) {
        std::vector<Node*> update(m_max_level, nullptr);
        Node* curr = m_head;

        for (int i = m_max_level - 1; i >= 0; --i) {
            while (curr->forwards[i] && curr->forwards[i]->key < elem) {
                curr = curr->forwards[i];
            }
            update[i] = curr;

            // 检查当前层是否已经存在相同的元素
            if (curr->forwards[i] && curr->forwards[i]->key == elem) {
                return; // 元素已存在，直接返回
            }
        }

        // 元素不存在，创建新节点并添加到跳表
        int level = randomLevel(); // 随机确定新节点的层数
        Node* newNode = new Node(elem, level);

        for (int i = 0; i < level; ++i) {
            newNode->forwards[i] = update[i]->forwards[i];
            update[i]->forwards[i] = newNode;
        }
        ++m_sz;
    }

    bool erase(ElemType elem) {
        std::vector<Node*> update(m_max_level, nullptr);
        Node* curr = m_head;

        for (int i = m_max_level - 1; i >= 0; i--) {
            while (curr->forwards[i] && curr->forwards[i]->key < elem) {
                curr = curr->forwards[i];
            }
            update[i] = curr;
        }

        curr = curr->forwards[0];
        if (curr && curr->key == elem) {
            for (int i = 0; i < m_max_level; i++) {
                if (update[i]->forwards[i] != curr) {
                    break;
                }
                update[i]->forwards[i] = curr->forwards[i];
            }
            delete curr;
            --m_sz;
            return true;
        }
        return false;
    }

    size_t size() {return m_sz;}
    bool empty() {return m_sz == 0;}

public:
    // 打印跳表
    void print() {
        for (int i = m_max_level - 1; i >= 0; i--) {
            Node* curr = m_head->forwards[i];
            if(curr) {
                std::cout << "Level " << i << ": ";
                while (curr) {
                    std::cout << curr->key << " ";
                    curr = curr->forwards[i];
                }
                std::cout << std::endl;
            }
        }
    }
};

}

#endif //SYNC_SET_HPP
